import os
import time
import json
import logging
from common.core.cron import CronBase
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.cache import RedisUtil
from common.enums.apply import VideoEnum
from common.enums.wallet import WalletEnum
from common.models.users import UserWalletModel
from common.models.video import VideoRecordsModel
from common.cache.queue_cache import QueueCache
from plugins.storage.driver import StorageDriver
from common.chain.video_server import VideoZhipuServer
from common.chain.video_server import VideoOpenAiHkServer

logger = logging.getLogger(__name__)


class Command(CronBase):
    @classmethod
    async def run(cls, **kwargs):
        try:
            # 读取任务
            record_id = await QueueCache.queue_pop(QueueCache.VIDEO_JOB)
            if not record_id:
                return

            # 查询任务
            record = await VideoRecordsModel.filter(id=record_id).first()
            if not record or record.status != VideoEnum.STATUS_ING:
                return

            # 处理响应
            if record.engine == "luma":
                await cls.handle_luma_check(record)
            elif record.engine == "zhipu":
                await cls.handle_zhipu_check(record)
            else:
                record.status = VideoEnum.STATUS_FAIL
                record.fail_reason = f"No processing channel {record.engine}"
                await record.save()
        except Exception as e:
            logger.error(f"handle_video error: {str(e)}")

    @classmethod
    async def handle_luma_check(cls, record: VideoRecordsModel):
        """
            Luma视频回调处理

            Document:
                https://www.volcengine.com/docs/6791/1366783

            Author:
                zero
            """
        cover_dir_path = ""
        video_dir_path = ""
        response = {}
        try:
            # 限流拦截
            lock_key: str = f"prevent:video_{str(record.id)}"
            lock_ack: bool = await RedisUtil.set(lock_key, str(record.id), 20, True)
            if not lock_ack:
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
                return False

            # 获取生成结果
            server = VideoOpenAiHkServer()
            response = await server.fetch(record.task_id)

            # 获取结果异常
            if response.get("error") or not response.get("data"):
                error = response.get("error") or {}
                await VideoRecordsModel.filter(id=record.id).update(
                    status=int(VideoEnum.STATUS_FAIL),
                    fail_reason=error.get("message", "未知错误"),
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI视频",
                    remarks=error.get("message", "未知错误")
                )
                return None

            # 处理生成结果
            data = response.get("data")
            if data[0].get("state") == "completed":
                # 成功
                result = data[0]["artifact"]
                cover_url: str = result["last_frame"]["url"]
                video_url: str = result["video"]["url"]

                # 文件扩展
                imageExt: str = str(cover_url.split(".")[-1]).strip().lower()
                videoExt: str = str(video_url.split(".")[-1]).strip().lower()

                # 存储路径
                filename: str = ToolsUtil.make_md5_str(str(video_url) + str(time.time()))
                curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
                cover_dir: str = f"storage/applies/video/{curr_now}/{filename}.{imageExt}"
                video_dir: str = f"storage/applies/video/{curr_now}/{filename}.{videoExt}"
                cover_dir_path: str = UrlUtil.to_root_path(path=cover_dir)
                video_dir_path: str = UrlUtil.to_root_path(path=video_dir)

                # 下载封面
                await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path)

                # 下载视频
                await ToolsUtil.download_file(url=video_url, save_path=video_dir_path)

                # 视频上传OSS
                engine: str = await UrlUtil.get_storage_engine()
                if engine != "local":
                    await StorageDriver.push(path=cover_dir_path, key=cover_dir)
                    await StorageDriver.push(path=video_dir_path, key=video_dir)

                # 更新记录
                await VideoRecordsModel.filter(id=record.id).update(
                    status=VideoEnum.STATUS_YES,
                    cover_url=str(cover_dir),
                    video_url=str(video_dir),
                    fail_reason="",
                    notify_snap=""
                )
            elif data[0].get("state") == "PROCESSING":
                # 进行中
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
        except Exception as e:
            if record.retry_count < 2:
                record.retry_count += 1
                record.start_time = int(time.time())
                await record.save()
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
            else:
                logger.error("Luma视频回调失败: [%s]", str(e))
                await VideoRecordsModel.filter(id=record.id).update(
                    status=VideoEnum.STATUS_FAIL,
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    fail_reason=str(e),
                    update_time=int(time.time())
                )
                if record.use_points > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record.user_id,
                        change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                        change_amount=record.use_points,
                        project="AI视频",
                        remarks=str(e)
                    )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if cover_dir_path and os.path.exists(cover_dir_path):
                    os.unlink(cover_dir_path)
                if video_dir_path and os.path.exists(video_dir_path):
                    os.unlink(video_dir_path)

    @classmethod
    async def handle_zhipu_check(cls, record: VideoRecordsModel):
        """
            智谱视频回调处理

            Document:
                https://www.volcengine.com/docs/6791/1366783
            """
        cover_dir_path = ""
        video_dir_path = ""
        response = {}
        try:
            # 限流拦截
            lock_key: str = f"prevent:video_{str(record.id)}"
            lock_ack: bool = await RedisUtil.set(lock_key, str(record.id), 20, True)
            if not lock_ack:
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
                return False

            # 获取生成结果
            server = VideoZhipuServer()
            response = await server.fetch(record.task_id)

            # 根据状态处理: [PROCESSING=处理中, SUCCESS=成功, FAIL=失败]
            if response.get("task_status") == "SUCCESS":
                result = response.get("video_result")[0]
                cover_url: str = result.get("cover_image_url")
                video_url: str = result.get("url")

                # 文件扩展
                imageExt: str = str(cover_url.split(".")[-1]).strip().lower()
                videoExt: str = str(video_url.split(".")[-1]).strip().lower()

                # 存储路径
                filename: str = ToolsUtil.make_md5_str(video_url)
                curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
                cover_dir: str = f"storage/applies/video/{curr_now}/{filename}.{imageExt}"
                video_dir: str = f"storage/applies/video/{curr_now}/{filename}.{videoExt}"
                cover_dir_path: str = UrlUtil.to_root_path(path=str(cover_dir))
                video_dir_path: str = UrlUtil.to_root_path(path=str(video_dir))

                # 下载封面
                await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path)

                # 下载视频
                await ToolsUtil.download_file(url=video_url, save_path=video_dir_path)

                # 视频上传OSS
                engine: str = await UrlUtil.get_storage_engine()
                if engine != str("local"):
                    await StorageDriver.push(path=cover_dir_path, key=cover_dir)
                    await StorageDriver.push(path=video_dir_path, key=video_dir)

                # 更新记录
                await VideoRecordsModel.filter(id=record.id).update(
                    cover_url=str(cover_dir),
                    video_url=str(video_dir),
                    status=VideoEnum.STATUS_YES
                )

                # 生成成功
                task_time = int(time.time()) - int(record.start_time)
                print(f"智谱视频成功: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s]")
            elif response.get("task_status") == "PROCESSING":
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
            else:
                error = response.get("error") or {}
                message = str(error.get("code", "0")) + ":" + error.get("message", "未知错误")
                await VideoRecordsModel.filter(id=record.id).update(
                    status=VideoEnum.STATUS_FAIL,
                    fail_reason=message,
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI视频",
                    remarks=message
                )
        except Exception as e:
            if record.retry_count < 2:
                record.retry_count += 1
                record.start_time = int(time.time())
                await record.save()
                await QueueCache.queue_push(QueueCache.VIDEO_JOB, record.id)
            else:
                logger.error("AI智谱视频回调失败: [%s]", str(e))
                await VideoRecordsModel.filter(id=record.id).update(
                    status=VideoEnum.STATUS_FAIL,
                    fail_reason=str(e),
                    notify_snap=json.dumps(response, ensure_ascii=False),
                    update_time=int(time.time())
                )
                if record.use_points > 0:
                    await UserWalletModel.inc(
                        scene=WalletEnum.POINTS,
                        user_id=record.user_id,
                        change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                        change_amount=record.use_points,
                        project="AI视频",
                        remarks=str(e)
                    )
        finally:
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                if cover_dir_path and os.path.exists(cover_dir_path):
                    os.unlink(cover_dir_path)
                if video_dir_path and os.path.exists(video_dir_path):
                    os.unlink(video_dir_path)
